package com.suo.RedBlackTree;

public class RedBlackTree {
    public static final int RED = 0;
    public static final int BLACK = 1;

    //节点
    private static class RBTreeNode {
        int color;
        Integer key;
        RBTreeNode left;
        RBTreeNode right;
        RBTreeNode parent;

        public RBTreeNode(Integer key) {
            this(key,null,null,null);
        }

        public RBTreeNode(Integer key, RBTreeNode left, RBTreeNode right, RBTreeNode parent) {
            this.key = key;
            this.left = left;
            this.right = right;
            this.parent = parent;
            this.color = BLACK;
        }
    }

    //根节点
    private RBTreeNode header;

    public RedBlackTree() {
        header = new RBTreeNode(null);
    }

    public void insert(Integer key) {
        insert(header,key);
    }

    public void printTree() {
        printTree(header);
    }

    //插入方法
    private void insert(RBTreeNode root, Integer key) {
        RBTreeNode node = new RBTreeNode(key);

        insert(root,node);
    }

    //插入方法
    private void insert(RBTreeNode root, RBTreeNode node) {
        RBTreeNode x = root;//跟着x周 用于指向父节点
        RBTreeNode y = null;//查询指针

        //1. 将红黑树当成一个正常的排序树插入
        while(x != null && x.key != null) {
            y = x;
            if (node.key < x.key) {
                x = x.left;
            }
            else {
                x = x.right;
            }
        }
        node.parent = y;
        //判断当前插入节点是父节点左孩子还是有孩子
        //当y == NULL 时说明此时x为根节点
        if (y == null) {
            header = node;
        }
        else {
            if (node.key < y.key) {
                y.left = node;
            }
            else {
                y.right = node;
            }
        }
        //2. 将出插入节点设置为红色
        node.color = RED;
        //3. 将二叉搜索树修正为红黑树
        fix(node);
    }

    /**
     * 红黑树的调整方法
     * 1. 如果插入节点的父节点为黑色则插入成功不需要调整
     * 2. 如果插入节点的父节点是红色
     *  2.1 如果叔叔节点为红色则将父节点与叔叔节点变为黑色，爷爷节点变为红色，在将爷爷节点作为子节点，循环执行直到根节点为止，最后再将根节点变为黑色
     *  2.2 如果叔叔节点为黑色则与AVL相同(LL,RR,LR,RL)但需要添加一步着色操作
     *   着色操作：
     *    2.2.1. 将插入节点x的父节点p和叔叔节点u的颜色变为黑色
     *    2.2.2. 将x的爷爷节点变为红色
     *    2.2.3. 将爷爷节点看作子节点，循环执行直到根节点为止，最后再将根节点变为黑色
     * @param node
     */
    private void fix(RBTreeNode node) {
        RBTreeNode parent = null;
        RBTreeNode grandParent = null;
        //如果我的父节点存在 并且父节点颜色时红色
        while((parent = node.parent) != null && parent.color == RED) {
            grandParent = parent.parent;
            //如果父节点时祖父节点的左孩子
            if (parent == grandParent.left) {
                //条件1 叔叔节点是红色
                RBTreeNode uncle = grandParent.right;
                if (uncle != null && uncle.color == RED) {
                    uncle.color = BLACK;
                    parent.color = BLACK;
                    grandParent.color = RED;
                    node = grandParent;
                    continue;
                }
                //条件2 叔叔节点是黑色 且当前节点是右孩子
                else if (parent.right == node) {
                    //LR先左旋后右旋
                    leftRotate(parent);
                    RBTreeNode temp = parent;
                    parent = node;
                    node = temp;
                }

                parent.color = BLACK;
                grandParent.color = RED;

                //条件3 叔叔是黑色的并且当前节点是左孩子
                rightRotate(grandParent);
            }
            //如果父节点是租房节点的右孩子
            else {
                RBTreeNode uncle = grandParent.left;
                if (uncle != null && uncle.color == RED) {
                    uncle.color = BLACK;
                    parent.color = BLACK;
                    grandParent.color = RED;
                    node = grandParent;
                    continue;
                }
                //如果是黑色 并且当前结点是左孩子
                if (parent.left == node) {
                    rightRotate(parent);
                    RBTreeNode temp = parent;
                    parent = node;
                    node = temp;
                }
                parent.color = BLACK;
                grandParent.color = RED;

                //如果叔叔是黑色 并且当前结点是右孩子
                leftRotate(grandParent);
            }
        }
        header.color = BLACK;
    }

    //左旋
    public void leftRotate(RBTreeNode p) {
        // 在当前节点不为null时，才进行左旋操作
        if (p != null) {
            // 先记录p的右儿子
            RBTreeNode rightChild = p.right;

            // 1. 空出右儿子的左子树
            p.right = rightChild.left;
            // 左子树不为空，需要更新父节点
            if (rightChild.left != null) {
                rightChild.left.parent = p;
            }

            // 2. 空出节点p的父节点
            rightChild.parent = p.parent;
            // 父节点指向右儿子
            if (p.parent == null) { // 右儿子成为新的根节点
                this.header = rightChild;
            } else if (p == p.parent.left) { // 右儿子成为父节点的左儿子
                p.parent.left = rightChild;
            } else { // 右儿子成为父节点的右儿子
                p.parent.right = rightChild;
            }

            // 3. 右儿子和节点p成功会师，节点p成为左子树
            rightChild.left = p;
            p.parent = rightChild;
        }
    }

    //右旋
    public void rightRotate(RBTreeNode p) {
        if (p != null) {
            // 记录p的左儿子
            RBTreeNode leftChild = p.left;

            // 1. 空出左儿子的右子树
            p.left = leftChild.right;
            // 右子树不为空，需要更新父节点
            if (leftChild.right != null) {
                leftChild.right.parent = p;
            }

            // 2. 空出节点p的父节点
            leftChild.parent = p.parent;
            // 父节点指向左儿子
            if (p.parent == null) { // 左儿子成为整棵树根节点
                this.header = leftChild;
            } else if (p.parent.left == p) { // 左儿子成为父节点左儿子
                p.parent.left = leftChild;
            } else { // 左儿子成为父节点的右儿子
                p.parent.right = leftChild;
            }

            // 3. 顺利会师
            leftChild.right = p;
            p.parent = leftChild;
        }
    }

    //遍历
    private void printTree(RBTreeNode root) {
        if (root == null) return;
        printTree(root.left);
        System.out.print(root.key + " ");
        printTree(root.right);
    }
}
